package com.wpr.shop.utils.servlet;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

/**
 * describe:
 *
 * @author zhanglisheng
 * @date 2018/10/25
 */
public class MyHttpResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream buffer = null;
    private ServletOutputStream out = null;
    private PrintWriter writer = null;

    public MyHttpResponseWrapper(HttpServletResponse resp) throws IOException {
        super(resp);
        this.buffer = new ByteArrayOutputStream();
        this.out = new MyHttpResponseWrapper.WapperedOutputStream(this.buffer);
        this.writer = new PrintWriter(new OutputStreamWriter(this.buffer, this.getCharacterEncoding()));
    }

    public ServletOutputStream getOutputStream() throws IOException {
        return this.out;
    }

    public PrintWriter getWriter() throws UnsupportedEncodingException {
        return this.writer;
    }

    public void flushBuffer() throws IOException {
        if (this.out != null) {
            this.out.flush();
        }

        if (this.writer != null) {
            this.writer.flush();
        }

    }

    public void reset() {
        this.buffer.reset();
    }

    public String getContent() throws IOException {
        this.flushBuffer();
        return new String(this.buffer.toByteArray(), "UTF-8");
    }

    private class WapperedOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream bos = null;

        public WapperedOutputStream(ByteArrayOutputStream stream) throws IOException {
            this.bos = stream;
        }

        public void write(int b) throws IOException {
            this.bos.write(b);
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener listener) {

        }
    }
}
